import numpy as np
import matplotlib.pyplot as plt

# 利用python编写如下程序：
# 假设某一润滑油厂样本数据集，其中包括训练集（train.txt文件）和测试集（test.txt文件）。数据集格式如下：
#
# Oil（含油量）	Trans（透明度）	Sti（粘度）	Oth（杂质量）	Weight（总重量）
# 34.43	5.32	7.22	0.31	243.23
#
# 请通过Python实现上述问题线性回归底层模型，并用此模型预测总重量，具体要求如下：
#
# 13　完成数据集的读取和特征缩放（6分）

# 数据集的读取
train = np.loadtxt(r'./x_oil_data/train.txt', delimiter=',')
test = np.loadtxt(r'./x_oil_data/test.txt', delimiter=',')
# 特征缩放
mu = train.mean(axis=0)
sigma = train.std(axis=0)
train -= mu
train /= sigma
mu = test.mean(axis=0)
sigma = test.std(axis=0)
test -= mu
test /= sigma
# 切分
x = train[:, :-1]
y = train[:, -1]
XX = np.c_[np.ones(len(x)), x]
x_test = test[:, :-1]
y_test = test[:, -1]
XX_test = np.c_[np.ones(len(x_test)), x_test]


# 14　实现代价函数（6分）
def model(x, theta):
    return x.dot(theta)


def cost_func(h, y):
    e = h - y
    sq = e ** 2
    j = sq.mean()
    j /= 2.0
    return j


# 15　实现梯度下降函数（6分）
def grad(x, y, alpha=0.001, num_iter=15000):
    m, n = x.shape
    group = num_iter // 20
    theta = np.zeros(n)
    j_his = np.zeros(num_iter)

    for i in range(num_iter):
        h = model(x, theta)
        e = h - y
        j = cost_func(h, y)
        j_his[i] = j
        if 0 == i % group:
            print(f'#{i + 1} cost = {j}')
        dt = 1.0 / m * x.T.dot(e)
        theta -= alpha * dt
    if 0 != i % group:
        print(f'#{i + 1} cost = {j}')
    return theta, j_his, h


# 16　完成测试集的数据预测，并计算在测试集上的代价函数值（5分）
theta, j_his, h_train = grad(XX, y)
print(f'Theta = {theta}')
plt.figure(figsize=[12, 4])
spr = 1
spc = 3
spn = 0
spn += 1
plt.subplot(spr, spc, spn)
plt.plot(j_his, label='cost function values')
plt.grid()
plt.legend()

# 17　以横轴为真实值，纵轴为预测值，画出测试集的对比散点图（5分）
spn += 1
plt.subplot(spr, spc, spn)
h_test = model(XX_test, theta)
plt.scatter(y_test, y_test, s=1, label='target values')
plt.scatter(y_test, h_test, s=1, label='hypothesis values')
plt.grid()
plt.legend()

spn += 1
plt.subplot(spr, spc, spn)
xx, yy = np.mgrid[-10:10:100j, -10:10:100j]
zz = np.zeros_like(xx)
for i, theta1 in enumerate(xx[:, 0]):
    for j, theta2 in enumerate(yy[0]):
        cost = cost_func(model(XX, np.r_[-3.26471308e-14, theta1, theta2, 3.97089840e-02, -1.01401654e-01].T), y)
        zz[i, j] = cost
plt.contour(xx, yy, zz)

plt.show()
